#include "../src/memory/memory_manager.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

// Test utility functions
static void test_print(const char* test_name, int success) {
    printf("[%s] %s\n", success ? "PASS" : "FAIL", test_name);
}

// Test memory manager initialization
static int test_mm_initialization() {
    printf("\n=== Testing Memory Manager Initialization ===\n");
    
    int success = 1;
    
    // Test with valid parameters
    memory_manager_t* mm = mm_init(64 * 1024 * 1024); // 64MB
    test_print("Memory Manager Initialization", mm != NULL);
    
    if (mm) {
        // Print initial statistics
        printf("\nInitial memory manager state:\n");
        mm_print_stats(mm);
        
        // Test cleanup
        int cleanup_result = mm_cleanup(mm);
        test_print("Memory Manager Cleanup", cleanup_result == 0);
        
        if (cleanup_result != 0) {
            success = 0;
        }
    } else {
        success = 0;
    }
    
    // Test with invalid parameters
    memory_manager_t* mm_null = mm_init(0);
    test_print("Invalid Parameter Handling", mm_null == NULL);
    
    return success;
}

// Test memory allocation
static int test_memory_allocation() {
    printf("\n=== Testing Memory Allocation ===\n");
    
    int success = 1;
    
    memory_manager_t* mm = mm_init(1024 * 1024); // 1MB
    if (!mm) {
        printf("Failed to initialize memory manager\n");
        return 0;
    }
    
    // Test basic allocation
    void* ptr1 = mm_alloc(mm, 1024, MEM_TYPE_PAGE_BUFFER);
    test_print("Basic Allocation", ptr1 != NULL);
    
    if (ptr1) {
        // Write some data to verify it's accessible
        memset(ptr1, 0xAA, 1024);
        bool data_ok = (((char*)ptr1)[0] == (char)0xAA && ((char*)ptr1)[1023] == (char)0xAA);
        test_print("Memory Write/Read", data_ok);
        
        if (!data_ok) {
            success = 0;
        }
    } else {
        success = 0;
    }
    
    // Test multiple allocations
    void* ptrs[10];
    int alloc_count = 0;
    
    for (int i = 0; i < 10; i++) {
        ptrs[i] = mm_alloc(mm, 512, MEM_TYPE_QUERY_BUFFER);
        if (ptrs[i]) {
            alloc_count++;
        }
    }
    
    bool multi_alloc_ok = (alloc_count >= 8); // Allow some to fail due to memory limits
    test_print("Multiple Allocations", multi_alloc_ok);
    
    if (!multi_alloc_ok) {
        success = 0;
    }
    
    // Test freeing memory
    if (ptr1) {
        mm_free(mm, ptr1);
    }
    
    for (int i = 0; i < alloc_count; i++) {
        mm_free(mm, ptrs[i]);
    }
    
    test_print("Memory Free", true); // If we get here without crashing, it worked
    
    // Test aligned allocation
    void* aligned_ptr = mm_aligned_alloc(mm, 1024, 64, MEM_TYPE_INDEX_BUFFER);
    bool aligned_ok = (aligned_ptr != NULL && ((uintptr_t)aligned_ptr % 64) == 0);
    test_print("Aligned Allocation", aligned_ok);
    
    if (aligned_ptr) {
        mm_free(mm, aligned_ptr);
    }
    
    if (!aligned_ok) {
        success = 0;
    }
    
    // Print final statistics
    printf("\nMemory allocation statistics:\n");
    mm_print_stats(mm);
    
    mm_cleanup(mm);
    return success;
}

// Test buffer pool operations
static int test_buffer_pool() {
    printf("\n=== Testing Buffer Pool ===\n");
    
    int success = 1;
    
    memory_manager_t* mm = mm_init(16 * 1024 * 1024); // 16MB
    if (!mm) {
        printf("Failed to initialize memory manager\n");
        return 0;
    }
    
    buffer_pool_t* bp = mm->buffer_pool;
    
    // Test getting pages
    buffer_page_t* page1 = bp_get_page(bp, 100);
    test_print("Get Page from Buffer Pool", page1 != NULL);
    
    if (page1) {
        // Test page pinning
        int pin_result = bp_pin_page(bp, page1);
        test_print("Pin Page", pin_result == 0);
        
        // Test marking dirty
        int dirty_result = bp_mark_dirty(bp, page1);
        test_print("Mark Page Dirty", dirty_result == 0);
        
        // Test page data access
        if (page1->page_data) {
            memset(page1->page_data, 0xBB, GUDB_PAGE_SIZE);
            bool data_ok = (((char*)page1->page_data)[0] == (char)0xBB);
            test_print("Page Data Access", data_ok);
            
            if (!data_ok) {
                success = 0;
            }
        }
        
        // Test unpinning
        int unpin_result = bp_unpin_page(bp, page1);
        test_print("Unpin Page", unpin_result == 0);
        
        // Test putting page back
        int put_result = bp_put_page(bp, page1);
        test_print("Put Page Back", put_result == 0);
        
        if (pin_result != 0 || dirty_result != 0 || unpin_result != 0 || put_result != 0) {
            success = 0;
        }
    } else {
        success = 0;
    }
    
    // Test cache behavior - get same page again (should be cache hit)
    buffer_page_t* page1_again = bp_get_page(bp, 100);
    bool cache_hit = (page1_again == page1);
    test_print("Cache Hit Behavior", cache_hit);
    
    if (!cache_hit) {
        success = 0;
    }
    
    // Test getting multiple pages
    buffer_page_t* pages[50];
    int page_count = 0;
    
    for (int i = 0; i < 50; i++) {
        pages[i] = bp_get_page(bp, 200 + i);
        if (pages[i]) {
            page_count++;
        }
    }
    
    bool multi_pages_ok = (page_count >= 30); // Should get most pages
    test_print("Multiple Page Gets", multi_pages_ok);
    
    if (!multi_pages_ok) {
        success = 0;
    }
    
    // Test page flushing
    int flush_result = bp_flush_all_dirty(bp);
    test_print("Flush All Dirty Pages", flush_result == 0);
    
    if (flush_result != 0) {
        success = 0;
    }
    
    // Print buffer pool statistics
    printf("\nBuffer pool statistics:\n");
    bp_print_stats(bp);
    
    mm_cleanup(mm);
    return success;
}

// Test LRU functionality
static int test_lru_functionality() {
    printf("\n=== Testing LRU Functionality ===\n");
    
    int success = 1;
    
    memory_manager_t* mm = mm_init(8 * 1024 * 1024); // 8MB - small buffer to force evictions
    if (!mm) {
        printf("Failed to initialize memory manager\n");
        return 0;
    }
    
    buffer_pool_t* bp = mm->buffer_pool;
    
    // Fill up the buffer pool
    const int num_pages = bp->total_pages + 10; // More than available
    buffer_page_t* pages[num_pages];
    
    for (int i = 0; i < num_pages; i++) {
        pages[i] = bp_get_page(bp, 1000 + i);
        // Access pattern: access first few pages multiple times
        if (i < 5) {
            bp_get_page(bp, 1000 + i); // Extra access to make them "hot"
        }
    }
    
    // Check that we have evictionss
    bool evictions_occurred = (bp->evictions > 0);
    test_print("LRU Evictions Occurred", evictions_occurred);
    
    if (!evictions_occurred) {
        success = 0;
    }
    
    // Check that frequently accessed pages are still in cache
    buffer_page_t* hot_page = bp_get_page(bp, 1000); // First page, should still be there
    bool hot_page_cached = (hot_page != NULL);
    test_print("Hot Pages Remain Cached", hot_page_cached);
    
    if (!hot_page_cached) {
        success = 0;
    }
    
    // Print final statistics
    printf("\nLRU test statistics:\n");
    bp_print_stats(bp);
    
    mm_cleanup(mm);
    return success;
}

// Test memory pressure handling
static int test_memory_pressure() {
    printf("\n=== Testing Memory Pressure Handling ===\n");
    
    int success = 1;
    
    // Create a small memory manager to easily trigger pressure
    memory_manager_t* mm = mm_init(512 * 1024); // 512KB
    if (!mm) {
        printf("Failed to initialize memory manager\n");
        return 0;
    }
    
    // Allocate memory until we run out
    void* ptrs[100];
    int alloc_count = 0;
    
    for (int i = 0; i < 100; i++) {
        ptrs[i] = mm_alloc(mm, 10 * 1024, MEM_TYPE_TEMP_BUFFER); // 10KB each
        if (ptrs[i]) {
            alloc_count++;
        } else {
            break; // Out of memory
        }
    }
    
    bool pressure_detected = (alloc_count < 100);
    test_print("Memory Pressure Detected", pressure_detected);
    
    if (!pressure_detected) {
        success = 0;
    }
    
    // Test pressure handling
    int pressure_result = mm_handle_pressure(mm);
    test_print("Memory Pressure Handling", pressure_result == 0);
    
    if (pressure_result != 0) {
        // This might fail if there's no memory to free, which is OK
    }
    
    // Try to allocate after pressure handling
    void* recovery_ptr = mm_alloc(mm, 1024, MEM_TYPE_TEMP_BUFFER);
    bool recovery_ok = (recovery_ptr != NULL);
    test_print("Post-Pressure Allocation", recovery_ok);
    
    // Note: This test might fail if pressure handling doesn't free enough memory
    // which is acceptable behavior
    
    // Free allocated memory
    for (int i = 0; i < alloc_count; i++) {
        if (ptrs[i]) {
            mm_free(mm, ptrs[i]);
        }
    }
    
    if (recovery_ptr) {
        mm_free(mm, recovery_ptr);
    }
    
    // Print final statistics
    printf("\nMemory pressure test statistics:\n");
    mm_print_stats(mm);
    
    mm_cleanup(mm);
    return success;
}

// Test concurrent access (basic)
static int test_thread_safety() {
    printf("\n=== Testing Thread Safety (Basic) ===\n");
    
    int success = 1;
    
    memory_manager_t* mm = mm_init(4 * 1024 * 1024); // 4MB
    if (!mm) {
        printf("Failed to initialize memory manager\n");
        return 0;
    }
    
    // Simple test: multiple allocations and frees
    // In a real test, we'd use multiple threads
    void* ptrs[20];
    
    for (int round = 0; round < 5; round++) {
        // Allocate
        for (int i = 0; i < 20; i++) {
            ptrs[i] = mm_alloc(mm, 1024 + i * 100, MEM_TYPE_QUERY_BUFFER);
        }
        
        // Free half
        for (int i = 0; i < 10; i++) {
            if (ptrs[i]) {
                mm_free(mm, ptrs[i]);
                ptrs[i] = NULL;
            }
        }
        
        // Allocate again
        for (int i = 0; i < 10; i++) {
            ptrs[i] = mm_alloc(mm, 2048, MEM_TYPE_INDEX_BUFFER);
        }
        
        // Free all
        for (int i = 0; i < 20; i++) {
            if (ptrs[i]) {
                mm_free(mm, ptrs[i]);
                ptrs[i] = NULL;
            }
        }
    }
    
    test_print("Multiple Alloc/Free Rounds", true);
    
    // Test buffer pool operations
    for (int i = 0; i < 100; i++) {
        buffer_page_t* page = bp_get_page(mm->buffer_pool, i % 50);
        if (page) {
            bp_mark_dirty(mm->buffer_pool, page);
            bp_put_page(mm->buffer_pool, page);
        }
    }
    
    test_print("Buffer Pool Operations", true);
    
    // Print final statistics
    printf("\nThread safety test statistics:\n");
    mm_print_stats(mm);
    
    mm_cleanup(mm);
    return success;
}

// Main test runner
int main() {
    printf("GuDB Memory Manager Tests\n");
    printf("=========================\n");
    
    int total_tests = 0;
    int passed_tests = 0;
    
    // Run all test suites
    struct {
        const char* name;
        int (*test_func)(void);
    } test_suites[] = {
        {"Memory Manager Initialization", test_mm_initialization},
        {"Memory Allocation", test_memory_allocation},
        {"Buffer Pool Operations", test_buffer_pool},
        {"LRU Functionality", test_lru_functionality},
        {"Memory Pressure Handling", test_memory_pressure},
        {"Thread Safety (Basic)", test_thread_safety}
    };
    
    int num_suites = sizeof(test_suites) / sizeof(test_suites[0]);
    
    for (int i = 0; i < num_suites; i++) {
        printf("\nRunning test suite: %s\n", test_suites[i].name);
        
        int suite_result = test_suites[i].test_func();
        total_tests++;
        
        if (suite_result) {
            passed_tests++;
            printf("Suite PASSED: %s\n", test_suites[i].name);
        } else {
            printf("Suite FAILED: %s\n", test_suites[i].name);
        }
    }
    
    // Print summary
    printf("\n========================================\n");
    printf("Test Results: %d/%d suites passed\n", passed_tests, total_tests);
    
    if (passed_tests == total_tests) {
        printf("All tests PASSED!\n");
        return 0;
    } else {
        printf("Some tests FAILED!\n");
        return 1;
    }
}